import json
import typer

from cli.src.cache.key import remove_bg_queue

from cli.src.cache import redis_client

from cli.src.db import repositories
from cli.src.db.schemas import OrderUpdate
from service.cli_service.ai import AIService


from config import appSettings

app = typer.Typer()


@app.command(name="predict")
def predict(model_dirname: str, image_path: str, bg_color: str):

    ai_service = AIService(model_dirname)
    """
    使用训练好的模型预测
    :param model_dirname:
    :param image_path:
    :param bg_color:
    :return:
    """
    print(model_dirname, image_path, bg_color)
    ai_service.replace_bg(
        model_path=model_dirname, image_path=image_path, bg_image_color=bg_color
    )
    return


@app.command(name="infer", help="ai模型部署")
def infer(folder: str):
    return


@app.command(name="remove_bg", help="ai去除背景")
def remove_bg():
    print(appSettings.rem_model_path)
    ai_service = AIService(appSettings.rem_model_path)
    red = redis_client.RedisClient().getRedis()
    while True:
        res = red.brpop(remove_bg_queue(), 3)
        # 处理数据
        # info_str = '{"user_id": 1,"image": "source/image/input/1-1.jpg","size": "xxx*xxx","pexil": "xxx*xxx","type": 1}'
        if res is None:
            continue
        if res[1] is None:
            continue
        info = json.loads(res[1])
        print(res[1])
        # 更新 订单状态
        out_pic = ai_service.remove_bg(info["image"])
        order = repositories.update_order(
            order_id=info["order_id"],
            order=OrderUpdate(out_pic=out_pic, status=2),
        )
        print(order.__dict__)
        print("处理完成")
        # TODO 处理完成发通知
    red.close()


if __name__ == "__main__":
    app()
